"""
数据源初始化脚本
初始化应急管理部数据源配置
"""
import datetime

from db.base import SessionLocal
from db.models.data_collect import DataSource, GlobalConfig


def init_emergency_data_sources():
    """初始化应急管理部数据源"""
    db = SessionLocal()
    try:
        # 应急管理部数据源配置
        emergency_sources = [
            {
                "source_id": "mem_bl",
                "name": "应急管理部-部令",
                "description": "应急管理部发布的部令文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/bl/",
                "category": "法规",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "bl",
                    "category_name": "部令"
                }
            },
            {
                "source_id": "mem_tb", 
                "name": "应急管理部-通报",
                "description": "应急管理部发布的通报文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/tb/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "tb",
                    "category_name": "通报"
                }
            },
            {
                "source_id": "mem_yjbgg",
                "name": "应急管理部-公告",
                "description": "应急管理部发布的公告文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/yjbgg/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "yjbgg",
                    "category_name": "公告"
                }
            },
            {
                "source_id": "mem_tz",
                "name": "应急管理部-通知",
                "description": "应急管理部发布的通知文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/tz/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "tz",
                    "category_name": "通知"
                }
            },
            {
                "source_id": "mem_h",
                "name": "应急管理部-函",
                "description": "应急管理部发布的函件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/h/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "h",
                    "category_name": "函"
                }
            },
            {
                "source_id": "mem_yj",
                "name": "应急管理部-意见",
                "description": "应急管理部发布的意见文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/yj/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "yj",
                    "category_name": "意见"
                }
            },
            {
                "source_id": "mem_qt",
                "name": "应急管理部-其他",
                "description": "应急管理部发布的其他文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/qt/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "qt",
                    "category_name": "其他"
                }
            },
            {
                "source_id": "mem_yjglbgg",
                "name": "应急管理部公报",
                "description": "应急管理部公报文件",
                "base_url": "https://www.mem.gov.cn/gk/tzgg/yjglbgg/",
                "category": "文件",
                "spider_config": {
                    "spider_type": "emergency",
                    "category": "yjglbgg",
                    "category_name": "应急管理部公报"
                }
            }
        ]
        
        # 交通部数据源配置
        mot_sources = [
            {
                "source_id": "mot_gov",
                "name": "交通运输部",
                "description": "交通运输部政府信息公开数据",
                "base_url": "https://xxgk.mot.gov.cn/2020/jigou/",
                "category": "法规",
                "spider_config": {
                    "spider_type": "mot_gov",
                    "category": "gov",
                    "category_name": "政府信息"
                }
            }
        ]
        
        # 合并所有数据源
        all_sources = emergency_sources + mot_sources
        
        # 插入或更新数据源
        for source_data in all_sources:
            existing = db.query(DataSource).filter(DataSource.source_id == source_data["source_id"]).first()
            if existing:
                # 更新现有数据源
                for key, value in source_data.items():
                    setattr(existing, key, value)
                existing.updated_at = datetime.datetime.utcnow()
            else:
                # 创建新数据源
                source = DataSource(**source_data)
                db.add(source)
        
        db.commit()
        print("应急管理部数据源初始化完成")
        
    except Exception as e:
        db.rollback()
        print(f"初始化数据源失败: {e}")
        raise
    finally:
        db.close()


def init_global_config():
    """初始化全局配置"""
    db = SessionLocal()
    try:
        # 默认全局配置
        default_configs = [
            {
                "config_key": "selected_sources",
                "config_value": ["mem_bl", "mem_tb", "mem_yjbgg", "mem_tz", "mot_gov"],  # 默认选中的数据源
                "description": "全局配置中选中的数据源列表"
            },
            {
                "config_key": "update_frequency", 
                "config_value": "1d",  # 默认一天更新一次
                "description": "数据更新频率：1d(一天)、3d(三天)、1w(一周)、1m(一个月)"
            },
            {
                "config_key": "enterprise_info",
                "config_value": "",  # 企业信息
                "description": "企业相关信息"
            }
        ]
        
        for config_data in default_configs:
            existing = db.query(GlobalConfig).filter(GlobalConfig.config_key == config_data["config_key"]).first()
            if existing:
                # 更新现有配置
                existing.config_value = config_data["config_value"]
                existing.description = config_data["description"]
                existing.updated_at = datetime.datetime.utcnow()
            else:
                # 创建新配置
                config = GlobalConfig(**config_data)
                db.add(config)
        
        db.commit()
        print("全局配置初始化完成")
        
    except Exception as e:
        db.rollback()
        print(f"初始化全局配置失败: {e}")
        raise
    finally:
        db.close()


def init_all():
    """初始化所有配置"""
    init_emergency_data_sources()
    init_global_config()
    print("数据源和全局配置初始化完成")


if __name__ == "__main__":
    init_all()
